'''
Description: 
Author: Egrt
Date: 2022-09-30 19:24:09
LastEditors: Egrt
LastEditTime: 2023-05-14 16:29:30
'''
import torch
import torch.backends.cudnn as cudnn
from nets.prnet import PRNet
from nets.arcface import Arcface
from utils.dataloader import LFWDataset
from utils.utils_metrics import test, plot_rocs

if __name__ == "__main__":
    #--------------------------------------#
    #   是否使用Cuda
    #   没有GPU可以设置成False
    #--------------------------------------#
    cuda            = True
    #--------------------------------------#
    #   主干特征提取网络的选择
    #   mobilefacenet
    #   mobilenetv1
    #   iresnet18
    #   iresnet34
    #   iresnet50
    #   iresnet100
    #   iresnet200
    #--------------------------------------#
    backbone        = "mobilefacenet"
    #--------------------------------------#
    #   输入图像大小
    #--------------------------------------#
    input_shape     = [112, 112, 3]
    #--------------------------------------#
    #   训练好的权值文件
    #--------------------------------------#
    model_path      = "model_data/ep134-loss4.233-val_loss5.954.pth"
    #--------------------------------------#
    #   LFW评估数据集的文件路径
    #   以及对应的txt文件
    #--------------------------------------#
    # lfw_dir_paths   = ["lfw", "mlfw", "dmlfw", "dmlfw2", "dmlfw3", "dmlfw4"]
    lfw_dir_paths   = ["cplfw"]
    lfw_pairs_path  = "model_data/pairs_CPLFW.txt"
    #--------------------------------------#
    #   评估的批次大小和记录间隔
    #--------------------------------------#
    batch_size      = 32
    log_interval    = 1
    #--------------------------------------#
    #   ROC图的保存路径
    #--------------------------------------#
    png_save_path   = "model_data/roc_test.png"
    tprs, fprs = [], []
    for lfw_dir_path in lfw_dir_paths:
        test_loader = torch.utils.data.DataLoader(
            LFWDataset(dir=lfw_dir_path, pairs_path=lfw_pairs_path, image_size=input_shape), batch_size=batch_size, shuffle=False)

        model = Arcface(backbone=backbone, mode="predict")

        print('Loading weights into state dict...')
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
        model  = model.eval()
        #------------------------------------------------------#
        #   载入PRNet模型及其权重
        #------------------------------------------------------#
        prnet = PRNet(3, 3)
        pretrained_dict = torch.load('model_data/prnet_pytorch.pth', map_location = device)
        prnet.load_state_dict(pretrained_dict)
        prnet = prnet.eval()
        

        if cuda:
            model = torch.nn.DataParallel(model)
            cudnn.benchmark = True
            model = model.cuda()

            prnet = torch.nn.DataParallel(prnet)
            prnet = prnet.cuda()

        tpr, fpr = test(test_loader, model, prnet, png_save_path, log_interval, batch_size, cuda)
        tprs.append(tpr)
        fprs.append(fpr)
    plot_rocs(fprs, tprs, png_save_path)
    
